import paramiko
import re
from typing import Optional, Tuple
from astrbot import Plugin, on_command, CommandContext, send_message, ask_question

# 全局状态（简化管理）
ssh_client: Optional[paramiko.SSHClient] = None
sftp_client: Optional[paramiko.SFTPClient] = None
current_session = {
    "connected": False,
    "username": "",
    "host": "",
    "port": 22
}

# 插件注册
class SimpleSSHPlugin(Plugin):
    name = "极简SSH插件"
    description = "Linux风格SSH工具，支持/ssh user@host -p端口、/ssh run、/ssh send"
    version = "2.0.0"

# ---------------------- 工具函数 ----------------------
def parse_ssh_command(cmd: str) -> Tuple[Optional[str], Optional[str], int]:
    """解析 /ssh user@host -p端口 格式"""
    # 匹配 用户名@主机 和 端口（可选）
    pattern = r"(\w+)@([\w\.\d]+)(?:\s+-p(\d+))?"
    match = re.search(pattern, cmd)
    if not match:
        return None, None, 22
    username = match.group(1)
    host = match.group(2)
    port = int(match.group(3)) if match.group(3) else 22
    return username, host, port

# ---------------------- 核心命令 ----------------------
@on_command("ssh", desc="连接SSH服务器（Linux风格）\n用法：/ssh user@server [-p端口]（默认22）")
async def ssh_connect(ctx: CommandContext):
    global ssh_client, current_session
    cmd = " ".join(ctx.args)
    
    # 解析参数
    username, host, port = parse_ssh_command(cmd)
    if not username or not host:
        await send_message(ctx, "❌ 格式错误！正确用法：/ssh user@server [-p端口]（如：/ssh root@192.168.1.100 -p5555）")
        return
    
    # 询问登录方式
    login_type = await ask_question(ctx, "请选择登录方式：\n1. 密码登录\n2. 密钥登录", timeout=30)
    if not login_type or login_type not in ["1", "2"]:
        await send_message(ctx, "❌ 选择无效，连接取消")
        return
    
    # 初始化SSH客户端
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        if login_type == "1":
            # 密码登录（隐式输入，避免明文）
            password = await ask_question(ctx, "请输入密码（输入后自动提交）", timeout=30, hidden=True)
            if not password:
                await send_message(ctx, "❌ 密码为空，连接取消")
                return
            ssh_client.connect(host, port, username, password, timeout=10)
        else:
            # 密钥登录（支持路径或直接输入密钥内容）
            key_input = await ask_question(ctx, "请输入密钥路径（如~/.ssh/id_rsa）或直接粘贴密钥内容", timeout=60)
            if not key_input:
                await send_message(ctx, "❌ 密钥为空，连接取消")
                return
            # 尝试解析为文件路径或密钥内容
            try:
                private_key = paramiko.RSAKey.from_private_key_file(key_input.strip())
            except:
                private_key = paramiko.RSAKey.from_private_key_string(key_input.strip())
            ssh_client.connect(host, port, username, pkey=private_key, timeout=10)
        
        # 连接成功更新状态
        current_session = {
            "connected": True,
            "username": username,
            "host": host,
            "port": port
        }
        await send_message(ctx, f"✅ 连接成功！\n{username}@{host}:{port}")
    except Exception as e:
        current_session["connected"] = False
        await send_message(ctx, f"❌ 连接失败：{str(e)}")

@on_command("ssh run", desc="执行SSH命令（实时输出）\n用法：/ssh run '命令'（如：/ssh run 'ls -l /root'）")
async def ssh_run(ctx: CommandContext):
    global ssh_client, current_session
    if not current_session["connected"] or not ssh_client:
        await send_message(ctx, "❌ 未连接服务器，请先使用 /ssh user@server 连接")
        return
    
    # 解析命令（支持带空格的命令，需用引号包裹）
    cmd = " ".join(ctx.args).strip()
    if not cmd or (cmd.startswith(("'", '"')) and cmd.endswith(("'", '"'))):
        cmd = cmd[1:-1]  # 去除引号
    if not cmd:
        await send_message(ctx, "❌ 命令为空！用法：/ssh run '命令'")
        return
    
    try:
        # 实时读取命令输出（逐行发送）
        await send_message(ctx, f"📝 执行命令：{cmd}\n--- 输出开始 ---")
        stdin, stdout, stderr = ssh_client.exec_command(cmd, get_pty=True)
        
        # 实时读取stdout
        for line in iter(stdout.readline, ''):
            line = line.strip()
            if line:
                await send_message(ctx, line)
        
        # 读取stderr（错误信息）
        err = stderr.read().decode("utf-8", errors="ignore").strip()
        if err:
            await send_message(ctx, f"❌ 错误信息：{err}")
        
        await send_message(ctx, "--- 命令执行结束 ---")
    except Exception as e:
        await send_message(ctx, f"❌ 命令执行失败：{str(e)}")

@on_command("ssh send", desc="发送文本/按键到远程终端\n用法：/ssh send '内容'（支持普通文本或按键，如'Enter'/'Ctrl+C'）")
async def ssh_send(ctx: CommandContext):
    global ssh_client, current_session
    if not current_session["connected"] or not ssh_client:
        await send_message(ctx, "❌ 未连接服务器，请先使用 /ssh user@server 连接")
        return
    
    content = " ".join(ctx.args).strip()
    if not content or (content.startswith(("'", '"')) and content.endswith(("'", '"'))):
        content = content[1:-1]
    if not content:
        await send_message(ctx, "❌ 发送内容为空！用法：/ssh send '文本或按键'")
        return
    
    try:
        # 发送文本（支持按键映射）
        stdin, _, _ = ssh_client.exec_command("cat", get_pty=True)  # 保持终端交互
        # 按键映射（常用按键）
        key_map = {
            "Enter": "\n",
            "Ctrl+C": "\x03",
            "Ctrl+D": "\x04",
            "Tab": "\t"
        }
        send_content = key_map.get(content, content)
        stdin.write(send_content + "\n")
        stdin.flush()
        await send_message(ctx, f"✅ 已发送：{'['+content+']' if content in key_map else content}")
    except Exception as e:
        await send_message(ctx, f"❌ 发送失败：{str(e)}")

@on_command("ssh exit", aliases=["ssh disconnect"], desc="断开SSH连接")
async def ssh_disconnect(ctx: CommandContext):
    global ssh_client, current_session
    if current_session["connected"] and ssh_client:
        try:
            ssh_client.close()
            current_session["connected"] = False
            await send_message(ctx, "✅ SSH连接已断开")
        except Exception as e:
            await send_message(ctx, f"❌ 断开失败：{str(e)}")
    else:
        await send_message(ctx, "❌ 当前未连接服务器")

# 插件入口
def setup():
    return SimpleSSHPlugin()
